import zmq
import cv2
import numpy as np
import json

def subscriber(address, topic):
    context = zmq.Context()
    socket = context.socket(zmq.SUB)
    socket.connect(address)
    socket.setsockopt_string(zmq.SUBSCRIBE, topic) # 订阅主题
    
    print("Waiting messages...")

    while True:
        try:
            topic, frame_bytes, result_bytes = socket.recv_multipart()
            print("Got message")

            # 解析topic
            topic = topic.decode()

            # 解析frame
            buffer = np.frombuffer(frame_bytes, dtype=np.uint8)
            frame = cv2.imdecode(buffer, cv2.IMREAD_COLOR)

            # 解析result，当内容为空时会受到""，则result会赋值为{}
            result = json.loads(result_bytes.decode('utf-8'))

            # 绘制出lines和circles
            if result != {}:
                lines = result["lines"] if "lines" in result else []
                circles = result["circles"] if "circles" in result else []

                # Draw 2D
                for line in lines:
                    cv2.line(frame, tuple(line[0]), tuple(line[1]), (255,255,0), 4, cv2.LINE_AA,)
                for circle in circles:
                    cv2.circle(frame, tuple(circle), 3, (0, 255, 255), -1, cv2.LINE_AA,)

            cv2.imshow('Frame', frame)
            key = cv2.waitKey(1)
            # escape = 27, use ESC to exit
            if key == 27:
                break
        except KeyboardInterrupt:
            print("Interrupted")
        except RuntimeError as e:
            print(e)

    cv2.destroyAllWindows()

if __name__ == "__main__":
    subscriber("tcp://localhost:5555", "hello")